import fitz

from receipt_identify.receipt.rects import GENERAL_RECEIPT_RECTS, VAT_RECEIPT_RECTS

doc1_path = "test\普通发票1.pdf"
doc2_path = "test\增值税发票2.pdf"


def outpur_png(page, output_path):
    # Save as png
    page.get_pixmap(dpi=300).save(output_path)


def draw_rect(page, rects):
    # Draw Rects
    for key, rect in rects.items():
        page.add_rect_annot(rect)


def main():
    with fitz.open(doc1_path) as doc1:
        page1 = doc1.load_page(0)
        draw_rect(page1, GENERAL_RECEIPT_RECTS)
        outpur_png(page1, "test/test1.png")

    with fitz.open(doc2_path) as doc2:
        page2 = doc2.load_page(0)
        draw_rect(page2, VAT_RECEIPT_RECTS)
        outpur_png(page2, "test/test2.png")


if __name__ == "__main__":
    main()
